/*
 * gmx_density.c
 *
 *  Created on: Apr 23, 2010
 *      Author: marchi
 */

#include "statutil.h"
#include "typedefs.h"
#include "smalloc.h"
#include "vec.h"
#include "copyrite.h"
#include "statutil.h"
#include "tpxio.h"
#include "fftgrid.h"
#include "index.h"
#include "pbc.h"
#include "rmpbc.h"

#include "get_mass.h"
//#include "centerofmass.h"
#include "xvgr.h"
#include "get_xplor.h"



rvec ngrid={1,1,1};
real*** Rho_sol;
real*** Rho_slt;
real*** Rho_sol2;
real*** Rho_slt2;
real*** Rho_tot;

real* mass;
enum { euSel,euRect, euTric, euCompact, euNR};

static void put_molecule_com_in_box(int unitcell_enum,int ecenter,
                                    t_block *mols,
                                    int natoms,t_atom atom[],
                                    int ePBC,matrix box,rvec x[])
{
    atom_id i,j;
    int     d;
    rvec    com,new_com,shift,box_center;
    real    m;
    double  mtot;
    t_pbc   pbc;
    calc_box_center(ecenter,box,box_center);
    set_pbc(&pbc,ePBC,box);
    if (mols->nr <= 0)
        gmx_fatal(FARGS,"There are no molecule descriptions. I need a tpr file for this pbc option.");
    for(i=0; (i<mols->nr); i++) {
        /* calc COM */
        clear_rvec(com);
        mtot = 0;
        for(j=mols->index[i]; (j<mols->index[i+1] && j<natoms); j++) {
            m = atom[j].m;
            for(d=0; d<DIM; d++)
                com[d] += m*x[j][d];
            mtot += m;
        }
        /* calculate final COM */
        svmul(1.0/mtot, com, com);
	for(d=0; d<DIM; d++)
        /* check if COM is outside box */
        copy_rvec(com,new_com);
        switch (unitcell_enum) {
        case euRect:
            put_atoms_in_box(box,1,&new_com);
            break;
        case euTric:
            put_atoms_in_triclinic_unitcell(ecenter,box,1,&new_com);
            break;
        case euCompact:
            put_atoms_in_compact_unitcell(ePBC,ecenter,box,1,&new_com);
            break;
        }
        if (debug) fprintf(debug," My debug %d %d \n",ePBC,ecenter);
        rvec_sub(new_com,com,shift);
        if (norm2(shift) > 0) {
            if (debug)
                fprintf (debug,"\nShifting position of molecule %d "
                         "by %8.3f  %8.3f  %8.3f\n", i+1, PR_VEC(shift));
            for(j=mols->index[i]; (j<mols->index[i+1] && j<natoms); j++) {
                rvec_inc(x[j],shift);
            }
        }else if(debug) fprintf (debug,"\n My Debug 2 %d  %f %f %f \n", i+1, new_com[XX], new_com[YY], new_com[ZZ]);
    }
}
static void put_residue_com_in_box(int unitcell_enum,int ecenter,
                                   int natoms,t_atom atom[],
                                   int ePBC,matrix box,rvec x[])
{
    atom_id i, j, res_start, res_end, res_nat;
    int     d, presnr;
    real    m;
    double  mtot;
    rvec    box_center,com,new_com,shift;
    calc_box_center(ecenter,box,box_center);
    presnr = NOTSET;
    res_start = 0;
    clear_rvec(com);
    mtot = 0;
    for(i=0; i<natoms+1; i++) {
#ifdef OLDVERSION
        if (i == natoms || (presnr != atom[i].resnr && presnr != NOTSET)) {
#else
        if (i == natoms || (presnr != atom[i].resind && presnr != NOTSET)) {
#endif
            /* calculate final COM */
            res_end = i;
            res_nat = res_end - res_start;
            svmul(1.0/mtot, com, com);
            /* check if COM is outside box */
            copy_rvec(com,new_com);
            switch (unitcell_enum) {
            case euRect:
                put_atoms_in_box(box,1,&new_com);
                break;
            case euTric:
                put_atoms_in_triclinic_unitcell(ecenter,box,1,&new_com);
                break;
            case euCompact:
                put_atoms_in_compact_unitcell(ePBC,ecenter,box,1,&new_com);
                break;
            }
            rvec_sub(new_com,com,shift);
            if (norm2(shift)) {
                if (debug)
#ifdef OLDVERSION
                    fprintf (debug,"\nShifting position of residue %d (atoms %u-%u) "
                             "by %g,%g,%g\n", atom[res_start].resnr+1,
#else
                    fprintf (debug,"\nShifting position of residue %d (atoms %u-%u) "
                             "by %g,%g,%g\n", atom[res_start].resind+1,
#endif
                             res_start+1, res_end+1, PR_VEC(shift));
                for(j=res_start; j<res_end; j++)
                    rvec_inc(x[j],shift);
            }
            clear_rvec(com);
            mtot = 0;
            /* remember start of new residue */
            res_start = i;
        }
        if (i < natoms) {
            /* calc COM */
            m = atom[i].m;
            for(d=0; d<DIM; d++)
                com[d] += m*x[i][d];
            mtot += m;
#ifdef OLDVERSION
            presnr = atom[i].resnr;
#else
            presnr = atom[i].resind;
#endif
        }
    }
}
static void center_x(int ecenter,rvec x[],matrix box,int n,int nc,atom_id ci[]){
    int i,m,ai;
    rvec cmin,cmax,box_center,dx;
    if (nc > 0) {
        copy_rvec(x[ci[0]],cmin);
        copy_rvec(x[ci[0]],cmax);
        for(i=0; i<nc; i++) {
            ai=ci[i];
            for(m=0; m<DIM; m++) {
                if (x[ai][m] < cmin[m])
                    cmin[m]=x[ai][m];
                else if (x[ai][m] > cmax[m])
                    cmax[m]=x[ai][m];
            }
        }
        calc_box_center(ecenter,box,box_center);
        for(m=0; m<DIM; m++)
            dx[m] = box_center[m]-(cmin[m]+cmax[m])*0.5;
        for(i=0; i<n; i++)
            rvec_inc(x[i],dx);
    }
}
static void acc_x(int ncent, int *cindex, rvec *x_in, rvec *x_out){
	for(int ii=0;ii<ncent;ii++){
		int i=cindex[ii];
		x_out[i][XX]+=x_in[i][XX];
		x_out[i][YY]+=x_in[i][YY];
		x_out[i][ZZ]+=x_in[i][ZZ];
	}
}
static void avg_x(int ncent, int *cindex,int n, rvec *x, rvec  add){
	for(int ii=0;ii<ncent;ii++){
		int i=cindex[ii];
		x[i][XX]/=(real) n;
		x[i][YY]/=(real) n;
		x[i][ZZ]/=(real) n;
		x[i][XX]+=add[XX];
		x[i][YY]+=add[YY];
		x[i][ZZ]+=add[ZZ];
	}
}

void gmx_density(int argc,char *argv[])
{
#ifdef OLDVERSION
  static char *desc[] = {
#else
  const char *desc[] = {
#endif
    "this is a small test program meant to serve as a template ",
    "when writing your own analysis tools. The advantage of ",
    "using gromacs for this is that you have access to all ",
    "information in the topology, and your program will be ",
    "able to handle all types of coordinates and trajectory ",
    "files supported by gromacs. Go ahead and try it! ",
    "This test version just writes the coordinates of an ",
    "arbitrary atom to standard out for each frame. You can ",
    "select which atom you want to examine with the -n argument."
  };
  FILE *fpxplor=NULL;
  t_topology top;
  int        ePBC;
  char       title[STRLEN], *grpnm;
  char       *grpnm2;
  t_trxframe fr;
  rvec       *xtop;
  matrix     box;
  int        status,status_out=-1;
  int        flags = TRX_READ_X;
  int        flagsb;
  int ncent, dcent;
  static rvec  newbox = {0,0,0}, shift = {0,0,0}, trans = {0,0,0};

#ifndef OLDVERSION
  output_env_t oenv;
#endif
  int nx,ny,nz;
  bool bIndex, bCenter=TRUE;
  atom_id      *cindex, *dindex;
  bool         bRmPBC,bPBCWhole,bPBCcomRes,bPBCcomMol,bPBCcomAtom,bPBC,bNoJump,bCluster;
  bool         bTrans=FALSE;
  bool 		bOXT=FALSE,bXplor=TRUE;
  const char  *warn;

  t_atoms      *atoms=NULL;
  rvec *x;

  int pbc_enum;
  enum
  {
      epSel,
      epNone,
      epComMol,
      epComRes,
      epComAtom,
      epNojump,
      epCluster,
      epWhole,
      epNR
  };
#ifdef OLDVERSION
  char *pbc_opt[epNR + 1] =
      { NULL, "none", "mol", "res", "atom", "nojump", "cluster", "whole",
          NULL };
  char *unitcell_opt[euNR+1] =
      { NULL, "rect", "tric", "compact", NULL };
  char *center_opt[ecNR+1] =
      { NULL, "tric", "rect", "zero", NULL };
#else
  const char *pbc_opt[epNR + 1] =
      { NULL, "none", "mol", "res", "atom", "nojump", "cluster", "whole",
          NULL };
  const char *unitcell_opt[euNR+1] =
      { NULL, "rect", "tric", "compact", NULL };
  const char *center_opt[ecNR+1] =
      { NULL, "tric", "rect", "zero", NULL };
#endif
  int unitcell_enum;
  enum
  { ecSel, ecTric, ecRect, ecZero, ecNR};
  int ecenter;
  t_pargs
		  pa[] ={
				  {"-grid", FALSE, etRVEC, {&ngrid}, "Grid size"},
                  { "-pbc", FALSE, etENUM,
                      { pbc_opt },
                      "PBC treatment (see help text for full description)" },
                  { "-ur", FALSE, etENUM,
                      { unitcell_opt }, "Unit-cell representation" },
                    { "-trans", FALSE, etRVEC,
                        { trans },
                        "All coordinates will be translated by trans. This "
                        "can advantageously be combined with -pbc mol -ur "
                        "compact." },
                  { "-boxcenter", FALSE, etENUM,
                      { center_opt }, "Center for -pbc and -center" },
  };
  t_filenm fnm[] = {
    { efTPS,  NULL,  NULL, ffREAD },   /* this is for the topology */
    { efTRX, "-f", NULL, ffREAD },      /* and this for the trajectory */
    { efTRX, "-oxt","coord.xtc", ffOPTWR },
    { efNDX, NULL,  NULL,         ffOPTRD },
  };
#define NFILE asize(fnm)
#define NPA asize(pa)

  CopyRight(stderr,argv[0]);
  /* This is the routine responsible for adding default options,
   * calling the X/motif interface, etc. */
#ifdef OLDVERSION
  parse_common_args(&argc,argv,PCA_CAN_TIME | PCA_CAN_VIEW,
		    NFILE,fnm,asize(pa),pa,asize(desc),desc,0,NULL);
#else
  parse_common_args(&argc,argv,PCA_CAN_TIME | PCA_CAN_VIEW,
		    NFILE,fnm,asize(pa),pa,asize(desc),desc,0,NULL, &oenv);
#endif
  /* We don't need any topology information to write the coordinates,
   * but to show how it works we start by writing the name and
   * charge of the selected atom. It returns a boolean telling us
   * whether the topology was found and could be read
   */
    pbc_enum   = nenum(pbc_opt);
    bPBCWhole  = pbc_enum==epWhole;
    bPBCcomRes = pbc_enum==epComRes;
    bPBCcomMol = pbc_enum==epComMol;
    bPBCcomAtom= pbc_enum==epComAtom;
    bNoJump    = pbc_enum==epNojump;
    bCluster   = pbc_enum==epCluster;
    bPBC       = pbc_enum!=epNone;
    bTrans    = opt2parg_bSet("-trans",NPA,pa);
    unitcell_enum = nenum(unitcell_opt);
    ecenter    = nenum(center_opt) - ecTric;
    /*
    int o;
    for(o=0;o<5;o++) printf("%s ",center_opt[o]);
    printf(" %d \n", nenum2(center_opt));
    printf(" ecenters %d %d \n",ecenter,ecTric);
    return;
     */
    bRmPBC = bPBCWhole || bPBCcomRes || bPBCcomMol;

    bOXT = opt2bSet("-oxt",NFILE,fnm);
	if (bOXT) {
		flagsb = flags | TRX_READ_X;
		status_out = open_trx(opt2fn("-oxt",NFILE,fnm),"w");
	}
    nx=(int) *ngrid;
    ny=(int) *(ngrid+1);
    nz=(int) *(ngrid+2);
    Rho_sol = mk_rgrid(nx, ny, nz);
    Rho_slt = mk_rgrid(nx, ny, nz);
    Rho_sol2 = mk_rgrid(nx, ny, nz);
    Rho_slt2 = mk_rgrid(nx, ny, nz);
    Rho_tot = mk_rgrid(nx, ny, nz);
	int i;
	real *ra=&Rho_sol[0][0][0];
	real *ra2=&Rho_sol2[0][0][0];
	real *rb=&Rho_slt[0][0][0];
	real *rb2=&Rho_slt2[0][0][0];
	for (i = 0; i < nx*ny*nz; i++) {
		ra[i] = 0.0;
		ra2[i]=0.0;
		rb[i] = 0.0;
		rb2[i] = 0.0;
	}
	bIndex = ftp2bSet(efNDX, NFILE, fnm);

	read_tps_conf(ftp2fn(efTPS, NFILE, fnm), title, &top, &ePBC, &xtop, NULL,
			box, TRUE);
	sfree(xtop);
	atoms = &top.atoms;
	if(bOXT){
		int natom=atoms->nr;
		x=calloc(natom,sizeof(rvec));
	}
	if (bCenter) {
		if (bIndex) {
			printf("Select group for centering\n");
			//	 rd_index(ftp2fn_null(efNDX, NFILE, fnm), 1, &ncent, &cindex, &grpnm);
			get_index(atoms,ftp2fn_null(efNDX, NFILE, fnm), 1, &ncent, &cindex, &grpnm);
			printf(" index!!  %d %d \n", ncent, cindex[ncent-1]);
		} else {
			gmx_fatal(FARGS, "Cannot do centering without Index file. Abort.");
		}
	}
	if (bIndex) {
		printf("Select group for density calculation\n");
		rd_index(ftp2fn_null(efNDX, NFILE, fnm), 1, &dcent, &dindex, &grpnm2);
			} else {
				gmx_fatal(FARGS, "Cannot do a density calculation without Index file. Abort.");
			}

	mass=malloc(sizeof(real)*top.atoms.nr);
	for(i=0;i<top.atoms.nr;i++) {
		mass[i]=top.atoms.atom[i].m;
	}
/* The first time we read data is a little special */
#ifdef OLDVERSION
	read_first_frame(&status,ftp2fn(efTRX,NFILE,fnm),&fr,flags);
#else
	read_first_frame(oenv, &status,ftp2fn(efTRX,NFILE,fnm),&fr,flags);
#endif

/* This is the main loop over frames */
	int n=0;
	real Volume=0.0;
	rvec angle_loc;
	rvec vec_cell={0.0, 0.0, 0.0}, angle_cell={0.0, 0.0, 0.0};

	do {
		/* coordinates are available in the vector fr.x
		 * you can find this and all other structures in
		 * the types directory under the gromacs include dir.
		 * Note how flags determines whether to read x/v/f!
		 */
		if (bCenter) center_x(ecenter,fr.x,fr.box,fr.natoms,ncent,cindex);
		t_pbc pbc;
		set_pbc(&pbc,ePBC,fr.box);

		if (bRmPBC)
			rm_pbc(&(top.idef),ePBC,fr.natoms,fr.box,fr.x,fr.x);

		if (bPBCcomAtom) {
			switch (unitcell_enum) {
				case euRect:
				put_atoms_in_box(fr.box,fr.natoms,fr.x);
				break;
				case euTric:
				put_atoms_in_triclinic_unitcell(ecenter,fr.box,fr.natoms,fr.x);
				break;
				case euCompact:
				warn = put_atoms_in_compact_unitcell(ePBC,ecenter,fr.box,
						fr.natoms,fr.x);
				break;
			}
		}
		if (bPBCcomRes) {
			put_residue_com_in_box(unitcell_enum,ecenter,
					fr.natoms,atoms->atom,ePBC,fr.box,fr.x);
		}
		if (bPBCcomMol) {
			put_molecule_com_in_box(unitcell_enum,ecenter,
					&top.mols,
					fr.natoms,atoms->atom,ePBC,fr.box,fr.x);
		}
		if(bTrans) {
			for(int i=0; i<fr.natoms; i++)
				rvec_inc(fr.x[i],trans);
		}

		if (bOXT) {
			acc_x(ncent, cindex, fr.x, x);
		}
		real dvol=det(fr.box);
		Volume+=dvol;
		dvol*=1000.0;
		dvol/=(real) (nx*ny*nz);

		cell_angles(fr.box,&angle_loc[0], &angle_loc[1], &angle_loc[2]);

		for(int d=0;d<DIM; d++){
			vec_cell[d]+=norm(box[d]);
			angle_cell[d]+=angle_loc[d];
		}

		if(get_mass(Rho_slt,Rho_sol,dvol,mass,nx,ny,nz,dcent,dindex, &fr)) gmx_fatal(FARGS,"Cannot average charges. Abort!\n");
		n++;
#ifdef OLDVERSION

	} while(read_next_frame(status,&fr));
#else
	} while(read_next_frame(oenv,status,&fr));
#endif

	Volume=Volume/(real) n;

	for(i=0;i<nx*ny*nz;i++){
		ra[i] /=(real) n;
		rb[i] /=(real) n;
	}

	matrix co;
	for(int d=0;d<DIM; d++){
		vec_cell[d]  /=(real) n;
		angle_cell[d]/=(real) n;
	}

	angles_cell(co, vec_cell,angle_cell);

	if(bXplor) {
		fpxplor=fopen("dens_slt.xplor","w+");
		get_xplor(Rho_slt,nx,ny,nz,co,fpxplor);
		ffclose(fpxplor);

		fpxplor=fopen("dens_sol.xplor","w+");
		get_xplor(Rho_sol,nx,ny,nz,co,fpxplor);
		ffclose(fpxplor);
	}

	if (bOXT) {
		rvec add;
		add[XX]=0.0;
		add[YY]=0.0;
		add[ZZ]=0.0;
		avg_x(ncent, cindex, n, x, add);

#ifdef OLDVERSION
		write_trx(status_out,ncent,cindex, &top.atoms,
              1,1.0,co,x,NULL);
#else
		write_trx(status_out,ncent,cindex, &top.atoms,
              1,1.0,co,x,NULL, NULL);
#endif
      }
	thanx(stderr);
}
